Add Unsloth Docker images (base + Studio) for any NVIDIA GPU host, Ampere through Blackwell#5748
Add Unsloth Docker images (base + Studio) for any NVIDIA GPU host, Ampere through Blackwell#5748danielhanchen wants to merge 120 commits into
Conversation
Adds a multi-stage Dockerfile producing an image that works on Ampere through Blackwell (sm_80 through sm_120: A100, RTX 30/40, H100, B100/B200, RTX 50-series, RTX 6000 Pro Blackwell). The build itself requires no GPU at all and runs on a free GitHub-hosted ubuntu-latest runner. How the GPU-less build works: 1. cu128 PyTorch wheels are fat binaries. torch._C._cuda_getArchFlags() returns 'sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120' regardless of which GPU compiled the image, because the wheels are cross-compiled upstream by the PyTorch team. 2. All deps resolve in a single uv pip install pass with explicit pins (torch==2.10.0, --extra-index-url cu128, no --torch-backend=auto, no install.sh). This prevents the silent cu cascade where bitsandbytes' transitive cuda-toolkit==13 dep upgrades torch to 2.12+cu130 in a later resolver pass, leaving xformers and other cu128 wheels stranded. 3. Build-time verification uses package metadata (importlib.metadata.version) and the raw torch._C._cuda_getArchFlags() accessor. We deliberately avoid import unsloth at build time because unsloth.__init__ calls torch.cuda.get_device_properties(0), which requires an actual CUDA device and is not bypassable. Import-time correctness is exercised at deploy time by smoke_test.py with --gpus all. 4. UNSLOTH_COMPILE_DISABLE=1 and CUDA_VISIBLE_DEVICES="" during the build stage prevent any code path from JIT-compiling kernels for the build host's compute capability and baking the resulting cache into the image. The deploy GPU produces its own cache on first use. Other notes: - --index-strategy unsafe-best-match is needed because the PyTorch wheel index serves an old requests==2.28.1 that conflicts with datasets>=2.32.2, which the default first-index-wins strategy rejects. - Extra is cu128-ampere-torch2100 (ampere precedes the torch version in the pyproject ordering). - No flash-attn in the base image. FA3 is hard-refused on Blackwell upstream and unsloth gracefully falls back to xformers + SDPA. Users on Ampere / Ada / Hopper who want FA2 can pip install flash-attn on top. - Two stages: nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 for the build, -cudnn-runtime for the deploy image. No nvcc in the published image. - A lockfile is emitted at /opt/unsloth-venv/requirements.lock.txt inside the image and can be extracted with docker/freeze.sh for byte-identical rebuilds even after PyPI moves on. CI workflow .github/workflows/docker-publish.yml: - Builds on ubuntu-latest on every push to main, every tag, weekly via cron, and manually via workflow_dispatch. Pushes to docker.io/unsloth/unsloth with cache via type=gha. - Optional smoke-test job runs on a self-hosted GPU runner if vars.HAS_GPU_RUNNER is set; skipped otherwise. End-to-end verification on sm_120 hardware is a nice-to-have, not a publish blocker. Validation: - Install path validated on a B200 host with CUDA_VISIBLE_DEVICES="" set (simulating the GPU-less CI runner): torch 2.10.0+cu128 holds, xformers 0.0.34, bitsandbytes 0.49.2, triton 3.6.0, transformers 5.5.0, trl 0.24.0, peft 0.19.1, accelerate 1.13.0. Arch flags include sm_100 and sm_120. - Runtime path validated end-to-end on B200: smoke_test.py imports unsloth, loads Llama-3.2-1B-Instruct-bnb-4bit in 4-bit, completes 5 LoRA steps with loss decreasing 4.11 -> 3.75. xformers fallback active as designed. Files: - docker/Dockerfile multi-stage cu128 build - docker/build.sh local build wrapper - docker/freeze.sh extract lockfile from a built image - docker/smoke_test.py runtime verification, run with --gpus all - docker/.dockerignore - .github/workflows/docker-publish.yml
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c6d92160f6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| docker pull ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest | ||
| docker run --rm --gpus all \ | ||
| ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest \ |
There was a problem hiding this comment.
Smoke-test the image built in this run, not
latest
The smoke-test job always pulls :latest, but this workflow also runs on tag pushes where latest is not guaranteed to be among the tags produced by metadata-action (it is only enabled on the default branch in this workflow). In that case, the smoke test can validate an older image and miss regressions in the freshly built tag from the current run.
Useful? React with 👍 / 👎.
| UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || 'main' }} | ||
| UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || 'main' }} |
There was a problem hiding this comment.
Build from triggering ref instead of hardcoding
main
For non-workflow_dispatch events (including tag pushes), github.event.inputs.* is unset, so these build args always resolve to main. That means images produced for v* tags can contain unsloth and unsloth-zoo code from main rather than the release ref that triggered the run, which breaks release correctness and reproducibility.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Code Review
This pull request introduces a Dockerized environment for Unsloth and unsloth-zoo, specifically optimized for NVIDIA Blackwell GPUs (sm_100 and sm_120). The changes include a multi-stage Dockerfile, build and freeze scripts, and a comprehensive smoke test to verify GPU compatibility and training functionality. Review feedback suggests optimizing the Dockerfile by removing a redundant installation of the uv tool, correcting a version mismatch for torchaudio to ensure consistency with the PyTorch stack, and relocating cache directories outside of the workspace to prevent issues when mounting host volumes at runtime.
| RUN ${VENV}/bin/pip install uv \ | ||
| && ${VENV}/bin/uv pip install \ | ||
| --python ${VENV}/bin/python \ |
There was a problem hiding this comment.
The uv tool is already installed at the system level in line 63. Installing it again inside the virtual environment at line 99 is redundant. Using the system-wide uv binary to install packages into the venv is more efficient and avoids unnecessary layers.
RUN uv pip install \
--python ${VENV}/bin/python \
| --python ${VENV}/bin/python \ | ||
| --index-strategy unsafe-best-match \ | ||
| --extra-index-url https://download.pytorch.org/whl/cu128 \ | ||
| "torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.11.0" \ |
There was a problem hiding this comment.
There appears to be a version mismatch for torchaudio. While torch is pinned to 2.10.0 and torchvision to 0.25.0 (which correctly follows the standard 0.(Y+15) mapping for Torch 2.10), torchaudio is set to 2.11.0. Typically, PyTorch and Torchaudio versions are released in sync (e.g., Torch 2.6.0 with Torchaudio 2.6.0). Using 2.10.0 ensures consistency across the stack.
"torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.10.0" \
| HF_HOME=/workspace/.cache/huggingface \ | ||
| TRITON_CACHE_DIR=/workspace/.cache/triton \ |
There was a problem hiding this comment.
Setting HF_HOME and TRITON_CACHE_DIR to subdirectories of /workspace (the WORKDIR) can lead to issues when users mount a host directory to /workspace at runtime. The mount will obscure the directories created during the build, forcing the application to recreate them at runtime, which can cause permission issues or redundant downloads. Moving these caches to a location outside of the workspace, such as /opt/cache, avoids these issues.
HF_HOME=/opt/cache/huggingface \
TRITON_CACHE_DIR=/opt/cache/triton \
| COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv | ||
|
|
||
| WORKDIR /workspace | ||
| RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR} |
There was a problem hiding this comment.
When someone launches the unsloth container, the common failure modes are not unsloth bugs -- they're Docker / nvidia-container-toolkit / driver issues that surface as cryptic CUDA errors deep in torch. The entrypoint catches the three that cover ~95% of "it doesn't work" reports up front: 1. nvidia-smi inside the container sees no GPU -> user forgot --gpus all, or host is missing nvidia-container-toolkit -> entrypoint prints the exact docker run flag and the toolkit install URL 2. nvidia-smi works but torch.cuda.is_available() is False -> host driver is older than CUDA 12.8 supports -> entrypoint prints the minimum driver version per architecture 3. compute capability < sm_80 -> entrypoint prints the supported architecture table and exits Each check fails with a clear, actionable message rather than a stack trace. Set UNSLOTH_SKIP_GPU_CHECK=1 to bypass (for docs builds, offline tooling, CI). run.sh wraps `docker run` with the flags people most often forget: --gpus all (without it, the new entrypoint refuses to start) --ipc=host (DataLoader workers need >64MB shm) --ulimit memlock=-1 (NCCL + CUDA pinned host buffers) --ulimit stack=64MB (some torch kernels OOM the default 8MB stack) Plus it mounts the host HF cache + Triton JIT cache so model downloads and compiled kernels persist across container runs, and forwards HF_TOKEN / WANDB_API_KEY / UNSLOTH_LICENSE only when they are set on the host. Usage: bash docker/run.sh # interactive python REPL bash docker/run.sh bash # shell in container bash docker/run.sh python /workspace/smoke_test.py bash docker/run.sh python /workspace/host/train.py # $PWD mounted at /workspace/host Verified locally: - No GPU visible: entrypoint refuses with driver-version message, exit 1 - B200 sm_100 visible: entrypoint prints GPU banner, exits cleanly into the user command (rc=0)
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 58693c4c73
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| [[ -n "${HF_TOKEN:-}" ]] && ENV_FORWARD+=(-e "HF_TOKEN=${HF_TOKEN}") | ||
| [[ -n "${WANDB_API_KEY:-}" ]] && ENV_FORWARD+=(-e "WANDB_API_KEY=${WANDB_API_KEY}") | ||
| [[ -n "${UNSLOTH_LICENSE:-}" ]] && ENV_FORWARD+=(-e "UNSLOTH_LICENSE=${UNSLOTH_LICENSE}") | ||
|
|
||
| set -x |
There was a problem hiding this comment.
Remove xtrace before invoking docker run with secrets
This script conditionally forwards HF_TOKEN, WANDB_API_KEY, and UNSLOTH_LICENSE, then enables set -x right before docker run, which prints the fully expanded command line. In any environment where those variables are set (local terminals with history/log capture or CI logs), their raw values are exposed in plaintext, creating an avoidable credential leak.
Useful? React with 👍 / 👎.
| import transformers | ||
|
|
||
| print(f"transformers {transformers.__version__}") |
There was a problem hiding this comment.
Import unsloth before transformers in smoke test
The default execution path runs training (check_tiny_train) but check_imports imports transformers first, which contradicts the later requirement that Unsloth be imported first for patched training behavior. This means the smoke test is not validating the intended Unsloth training path and can produce misleading pass/fail results (e.g., patches not applied or warnings/fallback paths during the LoRA step).
Useful? React with 👍 / 👎.
Single bash script that runs the full validation flow against the image:
1. Host pre-flight: docker version, nvidia-smi, nvidia-container-toolkit
runtime registered with docker.
2. Build the image (auto-detects the build context -- current dir,
docker/ subdir, or clones the docker-blackwell-build branch into
/tmp/unsloth-pr/).
3a. Smoke test: 5-step LoRA on Llama-3.2-1B-Instruct-bnb-4bit.
3b. Real workload: gpt-oss-20B fine-tuning notebook from
unslothai/notebooks, patched to max_steps=10, with the three
pre-train demo generations dropped for brevity. Auto-installs
triton_kernels at the SHA the upstream notebook pins for MXFP4.
All output is teed to /tmp/unsloth-docker-test/ (or --log-dir).
Usage:
bash docker/test_locally.sh # full run, ~15 min
bash docker/test_locally.sh --skip-notebook # blocks 1-3a only, ~3 min
bash docker/test_locally.sh --skip-build # reuse existing TAG
TAG=my:tag HF_TOKEN=hf_xxx bash docker/test_locally.sh
Each block fails fast with the exact log path to paste back.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: acbb16c8a1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| --python ${VENV}/bin/python \ | ||
| --index-strategy unsafe-best-match \ | ||
| --extra-index-url https://download.pytorch.org/whl/cu128 \ | ||
| "torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.11.0" \ |
There was a problem hiding this comment.
Align torchaudio pin with torch 2.10 stack
The install set mixes torch==2.10.0 with torchaudio==2.11.0, but TorchAudio wheels are release-coupled to specific Torch versions; this can make dependency resolution fail or force an unexpected Torch change, which breaks the Docker build’s stated guarantee that Torch stays on 2.10.0. This is especially risky here because the same layer also installs many transitive deps from multiple indexes, so one incompatible pin can fail the image build on CI.
Useful? React with 👍 / 👎.
| if ! command -v nvidia-smi >/dev/null 2>&1; then | ||
| err "nvidia-smi not found inside the container." | ||
| err "The CUDA runtime in this image is broken. Re-pull the image." | ||
| exit 1 |
There was a problem hiding this comment.
Avoid hard-failing when nvidia-smi binary is unavailable
Startup currently exits before any CUDA check if nvidia-smi is missing, but there are valid GPU runtimes (for example compute-only capability profiles) where CUDA is usable while nvidia-smi/NVML tools are not mounted. In those environments the container will refuse to start even though torch.cuda.is_available() could succeed, causing false negatives in production deployments that intentionally limit driver capabilities.
Useful? React with 👍 / 👎.
The Dockerfile uses BuildKit-only features (the # syntax=docker/dockerfile:1.7 parser directive and RUN ... <<'PY' heredocs added in dockerfile 1.3+). The legacy builder rejects the --progress flag at the CLI level and would fail later at the heredocs anyway. Detect docker buildx and use it when available (preserves --progress=plain output). Otherwise fall back to plain `docker build` with DOCKER_BUILDKIT=1 exported, which gets the BuildKit features without buildx's nicer formatting. Reproduces the failure path seen on Docker 28.2.2 without buildx installed: unknown flag: --progress ERROR docker build exited 125
Docker 28 removed the legacy image builder entirely. Setting
DOCKER_BUILDKIT=1 no longer falls back to a builtin builder -- it
delegates to buildx, which then errors out if buildx isn't installed:
ERROR: BuildKit is enabled but the buildx component is missing
or broken.
The Ubuntu docker.io package omits buildx by default, so users on
that path hit this immediately. Detect missing buildx up front and
print exact install commands for apt / dnf / manual binary instead
of attempting a fallback that cannot work.
If the user is not in the 'docker' group, every docker command after the pre-flight returns "permission denied while trying to connect to the Docker daemon socket at /var/run/docker.sock". This used to surface as a confusing buildx failure mid-Block-2, but the actual problem is a host permissions issue that's settable up front. Detect by running 'docker info' and checking its exit code (not just grep on its output -- a permission failure prints to stderr and returns non-zero, so the old grep-based check was a silent skip). Also clarify the nvidia-runtime WARN: on Docker 28+ with CDI mode this is a false positive most of the time. The real GPU-attach test is the smoke run in Block 3a, where the container entrypoint catches missing GPUs with an actionable message.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 23a5b43180
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| cd /workspace/host | ||
|
|
||
| echo "=== install triton_kernels (MXFP4 support for unsloth/gpt-oss-20b) ===" | ||
| pip install -q 'git+https://github.com/triton-lang/triton.git@0add68262ab0a2e33b84524346cb27cbb2787356#subdirectory=python/triton_kernels' 2>&1 | tail -5 |
There was a problem hiding this comment.
Propagate pip install failures through the tail pipeline
The generated run_notebook.sh script runs with set -e but not pipefail, and this line pipes pip install into tail -5; in Bash that means the pipeline exits with tail's status, so a failed pip install can be treated as success and the script continues until later steps fail with misleading errors (for example missing triton_kernels imports). This causes false-positive notebook validation and makes debugging CI/local failures much harder whenever installation fails (network hiccups, dependency conflicts, or bad commit hash).
Useful? React with 👍 / 👎.
Ubuntu 24.04 (noble) marks the system Python interpreter as externally-managed per PEP 668, so: curl get-pip.py | python python -m pip install -U pip uv fails inside the builder image with: error: externally-managed-environment This environment is externally managed The system-level pip and uv were never used: the very next RUN creates the venv at /opt/unsloth-venv, which bootstraps its own pip via the ensurepip module (provided by the python3.12-venv apt package). uv is then installed INTO the venv with the venv's pip, and used from there. Drop the two system-pip bootstrap lines. The venv path is unchanged. Reproduces on any Docker build of the unsloth-blackwell image against a noble base image (which our nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 is).
… / peft
unsloth_zoo/__init__.py guards against being imported standalone:
if "UNSLOTH_IS_PRESENT" not in os.environ:
raise ImportError("Please install Unsloth via `pip install unsloth`!")
The env var is set by unsloth/__init__.py at import time, so importing
unsloth must happen first. The old check_imports() imported xformers,
bnb, transformers, trl, peft, then unsloth_zoo -- which fired the guard
because unsloth had not been imported yet.
Reorder check_imports() to import unsloth (and unsloth_zoo) first, then
the rest. check_unsloth_import() becomes a thin re-import to keep the
"FastLanguageModel reachable" banner in the output.
Same fix the unsloth README has been recommending for years: "import
unsloth at the top of your file, before transformers/trl/peft."
Triton's nvidia backend lazily JIT-compiles a small C extension
(CudaUtils, in triton/backends/nvidia/driver.py) on first GPU access.
Without a C compiler and Python headers in the runtime image, the
very first forward pass of any Unsloth model dies with:
RuntimeError: Failed to find C compiler.
Please specify via CC environment variable.
The builder stage has build-essential and python3.12-dev so this
worked during the build's verification step (no GPU = no Triton kernel
call = no C extension build). But the runtime stage stripped those
out for size, so the failure only surfaces when a real user runs
training inside the container.
Add gcc + g++ + python3.12-dev to the runtime stage. Increases the
runtime image by ~250MB, which is the cost of letting Triton JIT
correctly. Pre-compiling CudaUtils at build time would need a real
CUDA device (the constructor calls cuda runtime functions), so
shipping the toolchain is the right trade-off.
Smoke-test validation on a fresh deploy host (AWS B200, not the build host)End-to-end validated What was validated
Unsloth's own banner inside the container: Loss progressionBit-for-bit identical to the internal validation on a GCP B200 (different host, same image): Bugs caught and fixed during validationEach was a real defect surfaced only by running the image on a fresh host:
Full gpt-oss-20B fine-tuning notebook run still pending; will post follow-up. |
…ia.com/cuda/gpus TORCH_CUDA_ARCH_LIST now covers the full set of compute capabilities NVIDIA publishes on https://developer.nvidia.com/cuda/gpus for x86_64 hardware, from Turing onward: sm_75 Turing T4, RTX 20-series, Quadro RTX sm_80 Ampere DC A100, A30 sm_86 Ampere A40, RTX A6000, RTX 30-series sm_89 Ada L4, L40, L40S, RTX 40-series sm_90 Hopper H100, H200, GH200 sm_100 Blackwell DC B100, B200, GB200 sm_103 Blackwell DC B300, GB300 sm_120 Blackwell RTX 50-series, RTX PRO 6000 Blackwell sm_121 Blackwell GB10 (DGX Spark) with +PTX on the highest entry so future arch revisions can JIT. Setting TORCH_CUDA_ARCH_LIST only affects nvcc invocations for any source build the user adds on top of this image (e.g. flash-attn, a custom CUDA op). The prebuilt cu128 wheels already include SASS for sm_70/75/80/86/90/100/120 (verified at build time via torch._C._cuda_getArchFlags()). Ada (sm_89), B300 (sm_103) and DGX Spark (sm_121) GPUs run via JIT-PTX from the nearest available arch. Jetson archs (sm_87 Orin, sm_110 Thor) are intentionally NOT included -- they require aarch64 wheels and this image is linux/amd64 only. Also lower the entrypoint's compute-capability gate from sm_80 to sm_75. Turing GPUs work, with the caveat that bfloat16 is unavailable; the entrypoint prints a NOTE in that case so Unsloth's fp16 fallback isn't a surprise.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dde5170e7a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ("sm_120", "Blackwell", "RTX 50-series, RTX PRO 6000 Blackwell"), | ||
| ("sm_121", "Blackwell", "GB10 (DGX Spark)"), | ||
| ) | ||
| if major < 7 or (major == 7 and minor < 5): |
There was a problem hiding this comment.
Reject pre-Ampere GPUs in entrypoint check
The entrypoint claims to fail fast on GPUs older than Ampere, but this condition only rejects devices below sm_75, so Turing (sm_75) is allowed through. That contradicts both the script’s own header (Unsloth requires sm_80+) and smoke_test.py, which exits on any cap[0] < 8; on T4/RTX20 hosts the container starts successfully and then fails later during validation/workloads instead of giving the intended immediate, actionable error.
Useful? React with 👍 / 👎.
HF Hub does not act as a generic OCI registry for arbitrary Docker images -- the registry.hf.space endpoint only serves images that Spaces have built, not images pushed by `docker push`. So we cannot do `docker push huggingface.co/user/repo:tag` for an Unsloth image. For cross-host testing where we want one canonical place to pull from (and Docker Hub credentials are not yet configured), wrap the manual flow into push/pull-shaped commands: hf_push.sh: docker save | pigz | huggingface-cli upload hf_pull.sh: huggingface-cli download | gunzip | docker load This is approximation, not real OCI semantics -- every push uploads the full ~4 GB blob, no layer dedup, no manifest negotiation. Good for testing across A100 / H100 / RTX 6000 boxes; the real release should go through .github/workflows/docker-publish.yml to Docker Hub, which gets layer dedup, multi-arch manifest support, and standard `docker pull` UX for users. Usage: bash docker/hf_push.sh unsloth-blackwell:test danielhanchen/unsloth-blackwell-docker bash docker/hf_pull.sh danielhanchen/unsloth-blackwell-docker unsloth-blackwell-test.tar.gz unsloth-blackwell:test
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4bfb4b891a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| with: | ||
| images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} | ||
| tags: | | ||
| type=raw,value=latest,enable={{is_default_branch}} |
There was a problem hiding this comment.
Prevent workflow_dispatch builds from overwriting
latest
This tag rule publishes latest whenever the run is on the default branch, but workflow_dispatch also allows overriding UNSLOTH_REF/UNSLOTH_ZOO_REF to arbitrary refs. A manual test run from main can therefore push a non-release image as latest, which makes downstream users pull an unintended build. Gate latest to trusted release flows (or only when both refs are the defaults) to avoid accidental retagging.
Useful? React with 👍 / 👎.
| set -euo pipefail | ||
|
|
||
| REPO="${1:?usage: hf_pull.sh <hf_repo> [<blob>] [<verify_tag>]}" | ||
| BLOB="${2:-unsloth-blackwell.tar.gz}" |
There was a problem hiding this comment.
Match
hf_pull default blob name to push naming
hf_push.sh uploads archives as <image-name>-<tag>.tar.gz, but hf_pull.sh defaults to unsloth-blackwell.tar.gz. If users run the documented short form bash docker/hf_pull.sh <hf_repo> after a normal push, the download will target a filename that was never uploaded and fail. Keeping both scripts on the same default naming convention avoids this broken default path.
Useful? React with 👍 / 👎.
…face-cli` In huggingface_hub >= 0.27 the `huggingface-cli` binary is deprecated and prints a "Use hf instead" notice then exits without doing the operation. The previous wrappers ran `huggingface-cli upload/download` silently, treated the deprecation exit as success, and uploaded nothing. Detect the new `hf` binary first and use that. If only the legacy `huggingface-cli` is on PATH (older installs), fall back with a WARN so users know the failure mode if anything goes sideways. Also: hf_pull.sh now asserts the downloaded file is non-empty (`test -s`) so we catch silent download failures before the `docker load` step.
…t call
`jupyter nbconvert --to script nb.ipynb --output nb 2>/dev/null` was
silently exiting 0 without producing the output file in some
environments (likely because jupyter/jupyter_core wasn't on PATH or
nbconvert's --output handling differed across versions). The 2>/dev/null
hid the underlying error, and `set -e` did not catch the missing-output
case because nbconvert itself returned 0.
Switch to a direct nbformat-based conversion:
pip install -q nbformat
python -c "import nbformat; nb=nbformat.read('nb.ipynb', as_version=4);
code='\n\n'.join(c.source for c in nb.cells if c.cell_type == 'code')
open('nb.py','w').write(code + '\n')"
Smaller dep set, no shell-out to a jupyter wrapper script, and an
explicit `test -s nb.py` afterwards catches any silent failure
before downstream steps try to read the file.
Reproduces the failure on RTX PRO 6000 Blackwell (sm_120, docker
29.2.1, ubuntu 24.04) where nbconvert's CLI silently no-op'd.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8344fa0a56
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| TRITON_CACHE_DIR=/workspace/.cache/triton \ | ||
| # Keep the arch list visible at runtime in case the user source-builds anything | ||
| # extra inside the container (e.g. a custom CUDA op). | ||
| TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0;10.0;12.0+PTX" |
There was a problem hiding this comment.
Include sm_75 in runtime arch list
The runtime image advertises support from Turing onward (see entrypoint.sh supported list), but TORCH_CUDA_ARCH_LIST here starts at 8.0. Any CUDA extension compiled inside the running container (the exact use case this env var comment describes) will be built without sm_75, so on T4/RTX20 hosts those kernels can fail at runtime with “no kernel image is available” despite the container claiming that architecture is supported.
Useful? React with 👍 / 👎.
The previous nbformat-based conversion dumped raw cell.source for every code cell. The gpt-oss-20B notebook's first cell uses Jupyter !shell magic to install dependencies: !pip install --upgrade -qqq uv !uv pip install -qqq ... \ git+https://github.com/triton-lang/triton.git@0add68... ... Dumped verbatim, the `@0add68...` token tripped the Python parser with "SyntaxError: invalid decimal literal" before training could even start. The container already has unsloth, triton, transformers, etc. baked in, so we don't need the notebook's install cell. Skip any cell whose source contains pip/install markers, and comment out stray !cmd / %magic lines in any other cells. Then assert nb.py parses with ast.parse() before trying to run it -- catches conversion failures up front instead of at training time. Reproduces on RTX PRO 6000 Blackwell (sm_120, fresh Docker 29.2.1 host) where the previous conversion produced an invalid nb.py.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 391532c031
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| git clone --depth 1 -b docker-blackwell-build \ | ||
| https://github.com/unslothai/unsloth.git /tmp/unsloth-pr 2>&1 | tail -3 |
There was a problem hiding this comment.
Clone stable ref in fallback build-context path
When this script is run outside the repo tree, the fallback path clones a hardcoded docker-blackwell-build branch. That branch name is PR-specific and can disappear after merge, so the fallback clone will fail and Block 2 cannot build at all. This breaks the script’s advertised “clone if needed” flow for users validating from a clean host; use a stable default ref (or a configurable ref input) instead of a transient PR branch.
Useful? React with 👍 / 👎.
Cross-host validation #2: sm_120 (RTX PRO 6000 Blackwell, 96GB) — full gpt-oss-20B fine-tuningEnd-to-end validated on a fresh GCP RTX PRO 6000 Blackwell Server Edition host. Image pulled from HF (huggingface.co/danielhanchen/unsloth-blackwell-docker), loaded into Docker, and exercised with both the smoke test AND the full gpt-oss-20B fine-tuning notebook. Setup
Smoke test (5-step LoRA on Llama-3.2-1B-bnb-4bit)``` Loss is ~0.4% off the B200 (sm_100) run (4.1108 → 3.7511); expected because sm_100 vs sm_120 Triton kernels produce slightly different bf16 rounding paths, deterministically. Full gpt-oss-20B fine-tuning (10 LoRA steps, MXFP4 + MoE)10 SFT steps on `HuggingFaceH4/Multilingual-Thinking` with the gpt-oss-20B MXFP4 model. Real workload, real MoE expert LoRA, real Harmony format inference at multiple reasoning_effort levels: ``` Unsloth: Detected MoE model with num_experts = 32 and target_modules = [...]. Trainable parameters = 3,981,312 of 20,918,738,496 (0.02% trained) step 1: loss=1.071 grad_norm=2.805 train_runtime = 141.7s (2.36 minutes) Post-train inference at `reasoning_effort="medium"` and `"high"` produced coherent French reasoning output via the Harmony format -- confirming MXFP4 weights are loading correctly, MoE expert routing works, and the Triton kernels JIT-compile for sm_120 at first use. What this validates
Additional bugs caught and fixed during this validation
|
Make the docker image multi-arch so DGX Spark (GB10, sm_121, aarch64) and
the Grace-Hopper / Grace-Blackwell SoCs (GH200 arm64, GB200 arm64) pull a
natively-built arm64 child from the same manifest. Runtime emulation is
NOT involved -- QEMU is used only for the cross-compile step on x86_64
CI runners; consumers on aarch64 hosts get a normal arm64 image and CUDA
works as on any other host.
Dockerfile:
* ARG TARGETARCH; switch unsloth extras between cu128-ampere-torch2100
(amd64, with xformers) and huggingface (arm64, no xformers -- there
is no cu128 aarch64 xformers wheel as of 0.0.34, so we fall back to
Unsloth's native SDPA path; ~5-10% slowdown but functionally complete).
* Build-time torch._C._cuda_getArchFlags() assertion: amd64 still
requires sm_120, arm64 accepts sm_120 or sm_121.
* Same TORCH_CUDA_ARCH_LIST on both arches; nvcc emits whatever's listed.
docker/setup_qemu.sh (new):
One-time host setup -- registers binfmt_misc handlers via
tonistiigi/binfmt and creates a 'unsloth-multiarch' docker-container
buildx builder. Required only on x86_64 build hosts targeting arm64.
docker/test_locally.sh:
--platform amd64|arm64 flag. Cross-builds verify QEMU is registered,
then build through the in-image arch-flags assertion. Smoke + notebook
blocks auto-skip when image arch != host arch (CUDA cannot run under
user-space QEMU + nvidia-container-toolkit cannot bridge a QEMU guest
to a real GPU).
.github/workflows/docker-publish.yml:
platforms: linux/amd64,linux/arm64 (single manifest, two children).
Timeout bumped 60 -> 150 min for the slower arm64-under-QEMU leg.
docker/setup-qemu-action@v3 with platforms: arm64 (was implicit before).
|
DGX Spark / linux/arm64 support added via QEMU at build time (e7cfcea). The image is now multi-arch: one Docker manifest with linux/amd64 + linux/arm64 children. What changed:
Arm64 GPU SoCs supported by the new variant:
End-to-end arm64 validation on actual hardware is pending -- the build-time assertion exercises the wheel resolution and the cu128 aarch64 fat binary, but the final proof is |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e7cfceadab
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| import unsloth_zoo | ||
|
|
||
| print(f"unsloth_zoo {unsloth_zoo.__version__}") | ||
| import xformers |
There was a problem hiding this comment.
Gate xformers import by platform in smoke test
The arm64 image path intentionally omits xformers (UNSLOTH_EXTRA="huggingface" in docker/Dockerfile) because no cu128 aarch64 wheel is expected, but check_imports() unconditionally does import xformers here. On arm64 runs this makes /workspace/smoke_test.py fail before training, so the published arm64 variant cannot pass the repository’s own runtime validation despite being a supported target.
Useful? React with 👍 / 👎.
| print(f"cuda build {torch.version.cuda}") | ||
| print(f"arches {arches}") | ||
| assert "sm_100" in arches, f"sm_100 missing: {arches}" | ||
| assert "sm_120" in arches, f"sm_120 missing: {arches}" |
There was a problem hiding this comment.
Accept sm_121 in smoke-test arch validation
This assertion hard-requires sm_120, but the same commit’s Dockerfile build validation explicitly treats arm64 Blackwell as valid when either sm_120 or sm_121 is present. As written, a valid arm64 build that reports only sm_121 will fail the smoke test with a false negative, even though the image is intended to support GB10/DGX Spark.
Useful? React with 👍 / 👎.
Sync the docker image branch with main (15 commits) so the branch's Python files carry main's current formatting and pre-commit.ci runs cleanly on a non-stale checkout.
Post-merge ruff-format-with-kwargs pass (the pre-commit.ci hook) on the merged loader.py; whitespace only, no logic change.
…6681) * docker: Colab-grade JupyterLab and Studio UX for the Blackwell image Stacks a Colab-like JupyterLab and Studio experience on top of the existing Blackwell image. Additive only: the training stack, CUDA/torch pinning, and the Studio/JupyterLab/sshd service trio are unchanged. JupyterLab labextension (prebuilt in a throwaway builder stage, so the runtime image stays Node-free): - Unsloth Dark (Monokai) theme, adaptive light/dark by system preference - Colab-style ArrowDown/Up cell navigation - top-bar Unsloth logo (stock Jupyter logo disabled and locked) - #@title lines render as collapsible Heading-2 form bars - Ctrl+A in a cell output selects only that output, not the whole notebook (the old behaviour ran notebook:select-all and was laggy) - right activity bar hidden by default - overrides.json: per-cell run button without auto-advance, labeled Restart and Run All, windowing off so collapsing an output does not snap to the cell top, news/update prompts suppressed Studio and login branding: Unsloth favicon, page logo, and a dark Unsloth login page that rotates through the curated Studio sloth stickers (fail-soft to the logo). Notebook organization and Colab compatibility (base image): - categorized folder view built from relative symlinks mirroring the README sections, rebuilt each boot; real .ipynb files never moved, and the symlink tree is invisible to the sync state machine - AMD-* notebooks shown only on an AMD/HIP host (autodetected) - Docker-only strip of the Colab "Run all on Colab" intro sentence from unedited notebooks (upstream notebooks unchanged) - hoist %%capture above a leading #@title form so the cell runs - the per-cell transformers-sidecar log is silent unless UNSLOTH_ENABLE_LOGGING=1 Dependency pinning and naming: the curated notebook extras are pinned to their resolved versions for reproducible rebuilds; decord is split into its own fail-soft install (no aarch64 wheel). The lean base image is renamed from :base to :core. Adds tests/validate_studio_features.py, a static self-test for the labextension plugins, overrides keys, and branding wiring. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docker: address review feedback on the JupyterLab/Studio UX - unsloth_nb_view.py: rebuilding the categorized view no longer deletes user files. The view is also JupyterLab's landing dir, so a user may save real notebooks there; _clear_view now unlinks only the symlinks we own and removes only folders that end up empty, leaving regular files in place. It also tests islink before isdir, so a view that is itself a symlink to a directory is unlinked instead of being walked into (which would have wiped the symlink target). - studio_launch.sh: derive the landing URL and preferred_dir from UNSLOTH_NOTEBOOKS_VIEW_DIR / UNSLOTH_SKIP_NOTEBOOK_VIEW, the same env the sync script uses, instead of hard-coding /workspace/Unsloth Notebooks. A relocated or disabled view no longer opens JupyterLab on a missing folder; it falls back to the default /lab over /workspace. - Dockerfile.studio: the labext-builder stage now installs Node 20 from NodeSource. Ubuntu 24.04's distro nodejs is 18, below JupyterLab 4.6's declared Node >=20 engine. Node stays confined to the throwaway builder stage, so the runtime image is unchanged. - .dockerignore: explicitly allowlist jupyter/install_sloth_stickers.py alongside its sibling jupyter assets, rather than relying on the directory re-inclusion. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docker: publish lean image as :core and full image as :studio Complete the base->core (and "studio as studio") tag rename so the publish workflow matches the user-facing helpers and the Dockerfile.studio header. - The lean training image now publishes as :core (core-<tag>, core-nightly, core-sha-*); run.sh / docker_confirm.* already told users to pull :core, but docker-publish.yml still tagged it :base, so that pull would have 404'd. The per-arch digest artifacts are renamed to match. - The full Studio image keeps :latest and gains a stable :studio alias, matching the Dockerfile.studio header. Both the merge and post-publish smoke-test metadata blocks are updated together. Internal "base image" wording (the layer Studio builds FROM) is left as-is. * docker: address second-round review feedback on the JupyterLab/Studio UX - studio_launch.sh: also gate the categorized-view landing URL on UNSLOTH_SKIP_NOTEBOOK_SYNC (the entrypoint skips building the view entirely in that mode), not just UNSLOTH_SKIP_NOTEBOOK_VIEW, so a no-sync container does not land on a missing folder. - Dockerfile.studio: scope the sticker-install "|| echo" fallback to only the sticker step via a { ...; } group. It was attached to the whole branding && chain, so a failure in a REQUIRED step (JS resolve, favicon/logo/login copy) was swallowed and the build continued with broken branding. - unsloth_nb_view.py: when creating the categorized symlinks, only replace our own stale symlinks; if a real user file already occupies that name, keep it and skip the link instead of os.remove-ing it. - overrides.json: drop doNotDisturbMode (it silenced ALL JupyterLab toasts, including kernel-restart / connection-drop feedback). The news/update prompts are already off via fetchNews / checkForUpdates. - Dockerfile: keep decord mandatory on amd64 (fail the build on a missing or incompatible wheel) and only fail-soft on arm64/other arches that have no wheel. - cellNav.ts: do not hijack ArrowUp/Down when focus is in an interactive output widget / form control, or while a completion popup is open, so ipywidgets controls and autocomplete at cell boundaries keep working. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docker: keep Studio branding RUN free of comments inside the line continuation Move the sloth-sticker fail-soft explanation above the RUN so no comment line sits between backslash-continued commands. BuildKit strips such comments, but keeping the RUN body a plain && chain removes the ambiguity for non-BuildKit builders and static linters. The { ...; } fail-soft scoping is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docker: AGPLv3 attribution + integrity guard for the Studio/JupyterLab image Make it obvious the image is built by Unsloth and hard to white-label out with a shallow find-and-replace, and surface the AGPLv3 license + copyright in the UI. Visible attribution (labextension): - Help > "About Unsloth Docker Studio" dialog (about.ts): Unsloth logo, the AGPLv3 notice, "Copyright 2026-Present the Unsloth team", and source/website/ license links. Added to the Help menu and the command palette. - The JupyterLab loading splash is replaced with a spinning Unsloth logo (splash.ts, provides ISplashScreen; honors prefers-reduced-motion). The stock @jupyterlab/apputils-extension:splash is disabled+locked at build time, like the stock logo. - AGPLv3 footer (license + copyright + links) on the branded login page. - Labextension relicensed AGPL-3.0-only; SPDX headers on every source file. Anti-tamper (no encoded/obfuscated strings -- plain readable text only; the one data URI is the logo image): - A canonical, plain-text attribution set lives in unsloth_branding.py with a TypeScript mirror (branding.ts) bundled verbatim into the labextension, so the phrase, copyright, links and plugin ids are spread across independent layers. - unsloth_branding.py verifies all of these across the installed files (AGPLv3 text, login footer, theme, labextension package + built bundle strings, logo, favicon) and fails loudly if any are missing. It runs at three layers: build time (fails the image build), the whole-container launcher (studio_launch.sh refuses to start), and as a jupyter_server extension (refuses to serve JupyterLab). - tests/studio/test_branding_guard.py: positive + per-marker negative coverage, plus a check that no base64/decoder obfuscation crept into the attribution. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docker: address #6681 review round 2 (colab magics, output select, branding guard) - unsloth_colab_compat.py: only hoist a leading `%%` cell magic above the Colab `#@title` form for magics whose body runs as code (capture/time/bash/python/ ...). Content magics (%%writefile, %%html, %%latex, ...) are left untouched so the form comment is never injected into the written file / rendered output. - outputSelect.ts: stop trusting the text selection anchor to decide ownership of Ctrl/Cmd+A. A stale selection inside an output survives a click onto a command-mode cell or the file browser, which made select-all keep re-selecting the old output. Gate on the keystroke target or the last pointer-down (reset to null on any click outside an output) instead. - unsloth_branding.py: also reject page_config.json that disables the Unsloth labextension or any of its plugin ids via disabledExtensions (dict or list form); that leaves the bundle on disk so the prior checks passed while the logo/About/splash attribution was stripped at load. Lock unsloth-jupyterlab in Dockerfile.studio as well (defense in depth), and add guard tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * labext: pin JupyterLab extension deps; confirm.ps1 /login probe Pin the unsloth-jupyterlab npm deps to exact versions matching the baked jupyterlab==4.6.0 (builder stays 4.5.9, its newest release) instead of floating ^/~ ranges, so the same commit always builds the same labextension bundle. Also probe JupyterLab /login (not /api, which 403s behind a password hash) in the Windows confirmation script. * docker: categorize AMD/domain notebooks and wire the feature validation into CI unsloth_nb_view.parse_readme only reset the folder section on level-3 (###) headings. The notebooks README carries level-1 domain headers (# AMD Notebooks, # Kaggle Notebooks) with their own nb/*.ipynb link tables and no intervening ###, so those notebooks were mis-filed under the previous stale section (all 148 AMD notebooks landed in Other Notebooks on an --amd build). Reset on any heading level and strip a leading emoji/symbol run so the domain notebooks get their own clean folder. Also run tests/validate_studio_features.py explicitly in the repo CPU job. It is named validate_* (not test_*) so pytest never collected it, which meant a regression in the notebook view, Colab compat, strip, JupyterLab defaults or login branding failed CI only when run by hand. * labext: use caret ranges so jlpm dedups JupyterLab/Lumino singletons The exact pins introduced earlier (@jupyterlab/* 4.6.0, @lumino/widgets 2.8.0, @jupyterlab/builder 4.5.9) break the Dockerfile.studio labext-builder stage. Exact-pinning the framework packages defeats jlpm's (yarn classic) hoisting: transitive @jupyterlab deps request caret ranges that resolve to newer patch releases (e.g. @jupyterlab/ notebook pulls @jupyterlab/cells ^4.6.0 -> a newer patch), so jlpm installs a second nested copy alongside the exact top-level one. Two copies of @jupyterlab/cells and @lumino/widgets in the tree produce TS2345 "not assignable" errors (protected-member/identity mismatch) and the build fails. Caret ranges let jlpm collapse every @jupyterlab and @Lumino package to a single hoisted copy, which is required for a JupyterLab prebuilt (federated) extension: at runtime those packages are shared singletons provided by the host JupyterLab, so the build-time versions only need to type-check against one consistent tree, not match an exact runtime patch. This is the version set the published image was built and validated with end to end. Verified by building the labext in isolation against the base image (Node 20 + bundled jlpm): caret ranges build clean (webpack compiled successfully); the exact pins fail with the duplicate-package TS errors. * ci(studio-backend): trigger on docker/** so the JupyterLab feature validation guards docker-only changes The 'Docker JupyterLab/notebook feature validation' step runs tests/validate_studio_features.py, which checks docker/jupyter (the labextension, overrides.json, login branding) and the docker notebook helpers. The pull_request paths filter listed studio/unsloth/tests but not docker/**, so a PR that only touches docker/ would skip that step and a regression in those files could pass CI. Add docker/** so the validation runs whenever the files it checks change. * jupyter: center the login card and place the attribution below it #site was a flex container using the default row direction with two children (the login card and the AGPLv3 attribution), so they rendered side by side: the card sat left of centre and the attribution floated up to the top-right. Stack them in a column so the card is horizontally centred and the attribution sits below it as a footer, matching the intended single-column layout. * jupyter: refresh Studio attribution, About dialog and loading splash - Attribution now reads 'Built by the Unsloth team' with a single Apache 2.0 / AGPLv3 license link (to the repo license section) on the login page and in the About dialog, replacing the plain 'Built by Unsloth. Licensed under the GNU AGPLv3.' line. The integrity guard, its canonical PHRASE and the branding tests are updated to match. - About dialog: left-align the link rows so the labels line up instead of each row centering independently; add an 'Unsloth Reference' link to the docs, and a Licenses section listing Unsloth Studio (AGPLv3) and Unsloth Core (Apache 2.0) alongside the full license link. - Loading splash now reads 'Loading Unsloth Docker' instead of the attribution label, via a dedicated SPLASH_LABEL constant. * docker: document the branding attribution as an AGPLv3 Section 7 notice Add docker/NOTICE and docker/jupyter/BRANDING.md so the Unsloth attribution that unsloth_branding.py enforces is also a written license condition, not only a build check. docker/NOTICE designates the attribution (the "Built by the Unsloth team" label, the copyright line, the license notice, the logo and theme, and the Help > About links) as required Appropriate Legal Notices under AGPLv3 Section 7(b), referencing /studio/LICENSE.AGPL-3.0 and /LICENSE. BRANDING.md is a human-readable note next to the guard describing what must stay, where it lives and how it is enforced. * ci(studio-backend): restore docker/** trigger path The docker/** pull_request path added in b558bc7 was dropped by a later rebase, so the "Docker JupyterLab/notebook feature validation" step (which runs tests/validate_studio_features.py against docker/jupyter branding and notebook helpers) no longer ran on PRs that only touch docker/. Re-add docker/** so a docker-only change is validated on the PR rather than only after merge to main. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
… to amd64
Base image (torch 2.11.0):
- amd64 unsloth extra: cu128-ampere-torch2100 -> cu128-ampere-torch2110.
The old extra pulls xformers 0.0.34, which hard-pins torch==2.10.0 and
conflicts with the torch==2.11.0 held throughout the build; the torch2110
family pulls xformers 0.0.35 (no torch pin) and resolves cleanly. This
needs an unsloth carrying the torch2110 CUDA extras on main, so merge the
torch2110 extras PR first (default UNSLOTH_REF=main).
- notebook-deps assertion: startswith('2.10.0') -> '2.11.0' so the layer
actually verifies the torch it now installs.
- refresh the torch2100/xformers 0.0.34 references in the surrounding
comments to the torch2110/0.0.35 line.
sm_103 (B300/GB300) JIT override (Codex item):
The cu13 NVRTC/ptxas override was arm64-only (sm_121), and its comment
claimed triton 3.6.0 bundles cu13 ptxas and set TRITON_PTXAS_PATH -- neither
was true: triton 3.6.0's bundled ptxas is CUDA 12.8 (V12.8.93, tops out at
sm_120) and TRITON_PTXAS_PATH was never set. So sm_103 (amd64) and even
sm_121 (arm64) Triton JIT were unfixed.
Run the cu13 install on both arches and actually wire the ptxas override:
- NVRTC swap (cu13 libnvrtc.so.13 over torch's bundled cu12.8 .so.12) now
runs on amd64 too.
- ENV TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas routes every Triton
JIT through the cu13 ptxas. Global rather than per-arch is safe: cu13.0
ptxas spans sm_70..sm_121 (verified: Volta/Turing/Ampere/Hopper through
Blackwell), so no regression for the older GPUs in the arch list.
Verified on amd64 in the built base image: cuda-nvrtc-13-0/cuda-nvcc-13-0
install cleanly from the base's CUDA repo, ptxas lands at
/usr/local/cuda-13.0/bin/ptxas (V13.0.88) and libnvrtc.so.13 at
/usr/local/cuda-13.0/lib64/. The sm_103/sm_121 runtime path itself is not
hardware-tested (no such GPU on hand); precompiled SASS still covers both
via sm_100/sm_120 forward-compat, so only JIT-heavy paths rely on this.
…before CUDA dedup
The Studio build symlinks the Studio venv's CUDA libs onto the base venv's
copies to reclaim ~3.7GB. That is only safe when both venvs run the same torch,
but the pre-dedup guard only checked the CUDA family (endswith('+cu128')). A
Studio venv that installed torch 2.10.0+cu128 (an installer capped below the
base's 2.11.0, or a build-time nvidia-smi fallback) would pass that check yet
mismatch the base's 2.11.0+cu128, and the dedup would link incompatible libs.
Capture the base venv's torch from its metadata and assert the Studio venv torch
equals it exactly (version and family) before the dedup runs, so a mismatch
fails the build loudly instead of silently linking skewed CUDA libs. Comparing
to the base venv also avoids hardcoding the version here. The Studio venv reaches
torch 2.11.0+cu128 via the installer's UNSLOTH_TORCH_INDEX_FAMILY=cu128 handling
and its CUDA torch spec allowing 2.11.x.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 386d3a7c74
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| CUDA_VERSION=12.8.1 | ||
| UBUNTU_VERSION=24.04 | ||
| PYTHON_VERSION=3.12 | ||
| UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || (startsWith(github.ref, 'refs/tags/') && github.ref_name) || github.sha || 'main' }} |
There was a problem hiding this comment.
Resolve UNSLOTH_REF before fan-out
When workflow_dispatch uses the default unsloth_ref=main or any branch name, this raw ref is fetched independently by each base matrix leg and again by the Studio build (UNSLOTH_STUDIO_REF uses the same expression). If that branch advances during the run, the amd64/arm64 children or the base and Studio layers can bake different Unsloth commits under one manifest/SHA tag; resolve the requested ref to one SHA in prepare and pass that output everywhere, like the zoo and notebooks refs already do.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d4dc8b6. Added a Resolve unsloth ref step to the prepare job that freezes the requested ref to one sha (same precedence as before: dispatch input, else pushed tag, else triggering sha, else main; ls-remote for a branch/tag), and rewired UNSLOTH_REF and UNSLOTH_STUDIO_REF to needs.prepare.outputs.unsloth_ref. Mirrors the existing zoo/notebooks resolution so all legs bake one commit.
| && if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \ | ||
| for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \ | ||
| if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \ | ||
| mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \ | ||
| ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \ |
There was a problem hiding this comment.
Swap the Studio NVRTC on amd64 too
On the full Studio image running on amd64 B300/GB300 (sm_103), this condition leaves the Studio venv's own nvidia/cuda_nvrtc/lib/libnvrtc.so.12 at CUDA 12.8. The base venv is swapped to cu13 and Triton gets a cu13 ptxas, but Studio's separate torch wheel can still load its bundled cu12.8 NVRTC via the wheel library path, so Studio-side NVRTC/Jiterator paths can fail on compute_103; apply the same symlink for amd64 or otherwise dedup/swap the Studio NVRTC there as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d4dc8b6. Dropped the arm64-only guard so the Studio venv NVRTC cu13 swap runs on both arches. amd64 sm_103 needs it too, the CUDA dedup never touches cuda_nvrtc, and the base cu13 layer installs cuda-nvrtc-13-0 on both arches so /usr/local/cuda-13.0/lib64/libnvrtc.so.13 exists regardless of TARGETARCH.
| "--upgrade-package", | ||
| "-P", |
There was a problem hiding this comment.
Filter uv upgrade-package values
When a notebook uses uv's upgrade selector for a protected package, e.g. uv pip install -P torch peft or --upgrade-package=transformers, this value is treated as an opaque option value and kept verbatim instead of passing through _KEEP / the transformers marker. uv documents this flag as allowing upgrades for a specific package, so installing a kept target can still refresh torch or transformers and clobber the baked stack; filter these values the same way as normal requirement specs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d4dc8b6. -P/--upgrade-package values now go through _KEEP: a protected value drops the flag+value pair (space and =forms), so uv pip install -P torch peft becomes uv pip install peft. Verified by running the shim on the exact example.
| _VALUE_FLAGS = { | ||
| "-r", | ||
| "--requirement", | ||
| "-c", | ||
| "--constraint", |
There was a problem hiding this comment.
Pair editable flags with their targets
Because -e/--editable are not treated as value-taking flags (checked python -m pip install --help: -e, --editable <path/url>), the scanner can leave a dangling editable flag after dropping a protected target. For example, pip install -e git+https://github.com/unslothai/unsloth.git#egg=unsloth peft drops the VCS URL but keeps -e, then executes pip install -e peft, failing instead of installing the kept package; handle editable as a paired target and drop both flag and value when the target is protected.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d4dc8b6. -e/--editable is now a value-taking flag paired with its target; a protected editable drops both flag and value, so pip install -e git+...unsloth...#egg=unsloth peft becomes pip install peft with no dangling -e. Non-protected editables (e.g. -e ./mypkg) are unchanged.
| if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")): | ||
| # A VCS / URL install can still name a protected package via the legacy | ||
| # `#egg=NAME` (or `&egg=NAME`) fragment, e.g. | ||
| # `git+https://github.com/unslothai/unsloth.git#egg=unsloth`. Pull that | ||
| # name out so _KEEP can drop it; otherwise the shim would exec the URL | ||
| # and reinstall a baked package into the venv. A non-protected egg name | ||
| # is returned too, but the caller keeps it as a normal target either way. | ||
| _egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token) | ||
| if _egg: | ||
| return _egg.group(1).lower().replace("_", "-") or None | ||
| return None # vcs / url / local path -> let it pass through |
There was a problem hiding this comment.
Parse protected wheel filenames before URL passthrough
This URL/path passthrough still lets bare wheel installs for protected packages bypass _KEEP: a valid cell such as pip install https://download.pytorch.org/.../torch-2.11.0+cu128-...whl returns None here, is kept as a positional target, and later runs the real pip, reinstalling torch into the baked venv. This is separate from the fixed PEP 508 and #egg= cases; parse wheel basenames (and requirement-file wheel lines) for protected distribution names before allowing URL/local-wheel passthrough.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d4dc8b6. _canon now parses the PEP 427 distribution name from a wheel URL/path, so pip install https://.../torch-2.11.0+cu128-...whl drops instead of reinstalling the baked torch. Non-protected wheels still pass through verbatim.
…p-shim edges) docker-publish.yml: freeze the requested unsloth ref to one sha in the prepare job before the matrix fans out. UNSLOTH_REF / UNSLOTH_STUDIO_REF were raw expressions re-evaluated per base arch leg and in the Studio build, so a mutable branch (the workflow_dispatch default unsloth_ref=main) advancing during the run could bake different unsloth commits under one manifest. Resolve once (same precedence: dispatch input, else pushed tag, else triggering sha, else main; ls-remote a branch/tag to a sha, mirroring the zoo/notebooks steps) and read needs.prepare.outputs.unsloth_ref everywhere. Dockerfile.studio: run the Studio venv NVRTC cu13 swap on both arches, not arm64 only. amd64 sm_103 (B300/GB300) needs cu13 NVRTC just as arm64 sm_121 does, and the CUDA dedup never touches cuda_nvrtc, so an amd64 Studio venv would otherwise keep its bundled cu12.8 libnvrtc and fail NVRTC/jiterator JIT on compute_103. The base cu13 layer installs cuda-nvrtc-13-0 on both arches, so the target .so.13 exists here regardless of TARGETARCH. unsloth_pip_shim.py: close three ways a protected package slipped past _KEEP. Treat -e/--editable as a value-taking flag paired with its target and drop both when the target is protected (was leaving a dangling -e that failed the cell); filter -P/--upgrade-package values through _KEEP (a named baked package could be refreshed while installing another target); and parse the PEP 427 distribution name out of a wheel URL/path so a bare `pip install https://.../torch-...whl` drops instead of reinstalling the baked torch. Non-protected editables, upgrade selectors, and wheels are unchanged. Adds tests/python/test_unsloth_pip_shim.py (18 regression tests, exec captured via a patched os.execv).
for more information, see https://pre-commit.ci
|
@codex review |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d6e559008f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| new_line, rewrote, inc_rec, inc_drp = _rewrite_include(line, stripped, src_dir, _depth) | ||
| out.append(new_line) |
There was a problem hiding this comment.
Filter editable entries in requirement files
Fresh evidence beyond the fixed command-line editable case: when a notebook runs pip install -r reqs.txt and reqs.txt contains a valid editable requirement like -e git+https://github.com/unslothai/unsloth.git#egg=unsloth, this branch treats every - line as an option/include and preserves non-include lines unchanged. python -m pip install --help documents -e, --editable <path/url>, and pip accepts that form from requirement files, so protected editable requirements can still bypass _KEEP and reinstall the baked stack; classify -e/--editable requirement-file lines before appending them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 251e3ed. An -e/--editable line inside a -r requirements file is now classified through _KEEP (new _parse_editable) and dropped when the target is protected, so -e git+...unsloth...#egg=unsloth in reqs.txt no longer reinstalls the baked stack. Non-protected editables are forwarded verbatim.
| # as an opaque option and a `-r` file is never filtered -- and worse, it | ||
| # never counts as a target, so a cell whose only target is that file | ||
| # silently no-ops and installs nothing. | ||
| if tok.startswith("--") and "=" in tok: |
There was a problem hiding this comment.
Handle attached short requirement options
For valid pip forms such as pip install -rrequirements.txt and pip install -cconstraints.txt peft (the local pip parser accepts -rfile/-cfile, and python -m pip install --help shows these options consume a <file>), this scanner only handles the separated form or long --flag=value. The attached short token starts with - and falls through as an opaque option, so an -r...-only cell no-ops with has_target false, while -c... stays unfiltered and can still apply protected pins during resolution.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 251e3ed. Attached short forms -rreqs.txt / -cconstraints.txt / -epath / -Pname are split into flag+value and routed through the same requirement-file / constraint / _KEEP handling as the separated form, so an attached -r no longer no-ops and -c/-e/-P values are filtered.
| if inc_rec and not recorded: | ||
| recorded = inc_rec |
There was a problem hiding this comment.
Do not record nested constraint pins as requests
Fresh evidence beyond the earlier constraint-file filtering fix: a requirements file that includes -c constraints.txt still propagates inc_rec here even when the nested include is a constraint, while the top-level -c path intentionally ignores _c_rec because a constraint is not a transformers install request. With reqs.txt containing -c constraints.txt plus a kept package, and the constraint pinning transformers==4.55.0, the shim writes the sidecar marker and the next model cell can activate that older transformers version even though it was only a resolver constraint; distinguish nested -c from nested -r before recording.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 251e3ed. A nested -c constraint include inside a -r file no longer propagates its transformers pin as a recorded request (a constraint is not an install request), mirroring the top-level -c path; the dropped-specs info still propagates.
| && git fetch -q --depth 1 origin "${UNSLOTH_STUDIO_REF}" \ | ||
| && git checkout -q FETCH_HEAD \ | ||
| && UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \ | ||
| UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \ |
There was a problem hiding this comment.
Let the Studio install resolve torch 2.11 on cu128
Fresh evidence beyond the earlier Studio torch-thread reply: this build now passes UNSLOTH_TORCH_INDEX_FAMILY=cu128, but the install.sh --local it invokes still keeps CUDA on TORCH_CONSTRAINT="torch>=2.4,<2.11.0" and only relaxes that constraint for */rocm7.2, not cu128. The Studio venv therefore resolves a 2.10.x+cu128 torch, and the equality check a few lines later compares it with the base venv's 2.11.0+cu128 and aborts every Studio image build; thread an exact/base torch spec or relax the CUDA cu128 constraint before running the installer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed, and this is provided by #6959 (branch studio-cuda-torch-2110): it adds a cu[0-9]*) case to install.sh widening cu128 to <2.12.0 (matched on _torch_index_leaf) plus the _CUDA_TORCH_PKG_SPEC repair-spec widening, so the Studio venv resolves 2.11.0+cu128 and the exact-match assert passes. The two must land together (#6959 with/before this PR).
| # above (triton 3.6.0's own ptxas is cu12.8, no sm_103/sm_121). Set globally, | ||
| # not per-arch: cu13.0 ptxas spans sm_70..sm_121 so it is correct for every GPU | ||
| # this image supports, and ENV cannot be made conditional per arch. | ||
| ENV TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas |
There was a problem hiding this comment.
Gate CUDA 13 JIT tools by driver/device
On otherwise-supported Ampere/Ada/Hopper/Turing hosts with 570–579 drivers, this global override sends every Triton JIT through CUDA 13 ptxas even though the entrypoint still advertises a 570+ driver floor except for sm_103/sm_121. NVIDIA's CUDA compatibility table lists CUDA 13.x as requiring a 580+ driver, so ordinary training JITs can fail on drivers the image currently accepts; enable the CUDA 13 ptxas/NVRTC swap only for sm_103/sm_121 or raise the documented/enforced driver floor.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 251e3ed. Confirmed: a cu13 cubin needs a >=580 driver to LOAD even when targeting an older arch (CUDA has forward, not backward, cross-major driver compat), so a global cu13 override breaks 570-579 hosts on Turing..sm_120. Removed the global ENV; select_cuda_jit_tools() in entrypoint.sh reads the device compute_cap and activates cu13 only for sm_103/sm_121 (which ship >=580 drivers), else keeps the bundled cu12.8 ptxas and restores the cu12.8 NVRTC in both venvs. 7 device-gating unit tests added.
…gate cu13 JIT tools) unsloth_pip_shim.py: close three more ways a protected package slipped past _KEEP. An editable line (-e/--editable <target>) inside a -r requirements file is a real install target, so a protected editable there is now classified and dropped like the command-line case (new _parse_editable). pip/uv accept the attached short forms -rreqs.txt / -cconstraints.txt / -epath / -Pname as one token; these were falling through as opaque options (so an attached -r-only cell no-op'd and an attached -c/-e/-P value bypassed _KEEP), so the 2-char flag is now split from its value and routed through the separated-form handling. And a nested -c constraint inside a -r file no longer records its transformers pin as an install request (a constraint is not a request; mirrors the top-level -c path). entrypoint.sh / Dockerfile: gate the CUDA 13 ptxas + NVRTC to sm_103 / sm_121 at runtime instead of a global build-time default. A cu13 cubin needs a >= 580 driver to LOAD even when it targets an older arch (CUDA has forward, not backward, cross-major driver compatibility), but the image supports Turing.. sm_120 on a 570+ driver, so the previous global TRITON_PTXAS_PATH ENV + cu13 NVRTC symlink would break ordinary Triton/NVRTC JIT on 570-579 driver hosts. The build still bakes cu13 (saving the cu12.8 NVRTC as .cu128.orig); a new select_cuda_jit_tools() in the entrypoint reads the device compute_cap and only activates cu13 for sm_103/sm_121 (which ship >= 580 drivers), otherwise leaving Triton on its bundled cu12.8 ptxas and restoring the cu12.8 NVRTC in both the base and Studio venvs. The base ENTRYPOINT runs for the Studio image too. Adds 9 pip-shim regression tests and tests/sh/test_select_cuda_jit_tools.sh (7 device-gating cases); registers the latter in CI and tests/run_all.sh.
|
@codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3bb40e47fe
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| ) | ||
| if _dref: | ||
| return _dref.group(1).lower().replace("_", "-") or None | ||
| if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")): |
There was a problem hiding this comment.
Parse bare wheel filenames before passthrough
The fresh bypass is a bare relative wheel filename: pip accepts pip install torch-...whl from the current directory, but this guard only enters the wheel-basename parser for URLs/VCS or paths starting with .//. That token falls through as the normalized full filename, misses _KEEP, and the real pip can reinstall torch/unsloth/etc. into the baked venv; run the .whl parser before this guard or include bare *.whl tokens.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b3649d4. _canon now parses a bare wheel filename (no ./ or / prefix and no scheme) via its PEP 427 basename before the plain-spec fallthrough, so pip install torch-2.11.0-...whl from the CWD resolves to torch and is dropped by _KEEP. Added test_bare_torch_wheel_filename_dropped / _in_subdir_dropped / _unprotected_kept.
| local cc="" nvrtc_dir orig | ||
| if command -v nvidia-smi >/dev/null 2>&1; then | ||
| cc="$( { nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true; } \ | ||
| | head -n1 | tr -d '[:space:]' )" |
There was a problem hiding this comment.
Check every visible GPU before restoring cu128 JIT tools
On mixed GPU hosts where nvidia-smi lists an H100/B200 before a B300/GB10, this line discards all but the first compute capability. Startup then takes the default branch, restores cu12.8 NVRTC and leaves Triton's ptxas unset even though a visible sm_103/sm_121 device still needs the CUDA 13 tools to JIT; workloads pinned to that later GPU hit the compute_103/121 failures this image is trying to avoid. Consider scanning all visible caps and enabling cu13 if any visible GPU requires it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b3649d4. select_cuda_jit_tools() now scans every visible compute_cap (a while-read over the full nvidia-smi list) and enables cu13 if ANY GPU is sm_103/sm_121, so a B300/GB10 sitting behind an H100/B200 still gets the cu13 tools. Those parts only ship on >=580 drivers, so the whole host tolerates cu13 cubins for every arch present. Added mixed-host tests (H100+B300, B200+GB10, and the order-reversed B300+H100).
| except OSError: | ||
| return path, None, [] # remote URL / unreadable -> let the real tool handle it |
There was a problem hiding this comment.
Reject remote requirement files before pip sees them
For a notebook cell like pip install -r https://.../reqs.txt, pip accepts the URL form (the local pip install --help says the requirement target can be a file or URL), but this path treats the URL as unreadable and passes it through unchanged. Protected specs inside that remote file never go through _KEEP or the transformers sidecar marker, so the real pip can still reinstall torch/transformers into the baked venv; either fetch/filter URL requirement files or reject them in notebook mode.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b3649d4. Remote (URL) -r/-c requirement/constraint files are now refused in shim mode: the top-level form in all three arg spellings (-r URL, --requirement=URL, -rURL) and nested -r/-c URL includes inside a local file, since their pins cannot be inspected before the real tool would fetch and install them. A remote -r as the sole target no-ops with no dangling flag; other targets on the same line still install. Added test_remote_requirement_url_* and test_nested_remote_include_dropped.
| dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-") | ||
| if dist: | ||
| return dist | ||
| return None # vcs / url / local path -> let it pass through |
There was a problem hiding this comment.
Drop protected VCS URLs without egg fragments
Fresh evidence beyond the fixed #egg= case: pip also accepts VCS project URLs without an egg fragment (local help lists VCS project URLs, and this repo recommends git+https://github.com/huggingface/transformers.git in unsloth/models/loader.py). Those tokens reach this fallback with _canon(...) == None, are kept as install targets, and can overwrite the baked transformers/unsloth stack without _KEEP or the marker ever firing; infer protected names from known repo basenames or block unclassified protected VCS URLs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b3649d4. _canon now infers the distribution from an egg-less VCS URL by its repo basename (git+https://github.com/huggingface/transformers.git -> transformers, unslothai/unsloth-zoo.git -> unsloth-zoo), stripping a trailing @ref and .git, so the egg-less form loader.py recommends is matched against _KEEP. A non-protected repo returns its basename and the caller keeps it as a normal target either way. Added test_vcs_url_without_egg_protected_dropped / _with_ref_dropped / _unprotected_kept.
| continue | ||
| name = _canon(tok) | ||
| if name is None: | ||
| keep_args.append(tok) # bare flag, or a positional url / path / vcs |
There was a problem hiding this comment.
Drop resolver-wide reinstall flags in notebook shim
In notebook cells such as pip install --force-reinstall peft or pip install --force-reinstall transformers==4.55.0 peft, this branch keeps the global reinstall flag after protected specs are dropped. Local pip install --help defines --force-reinstall as reinstalling all packages even if up-to-date, so the real pip can still replace protected dependencies pulled by the kept target (for example transformers/torch) even though direct protected requirements were filtered; strip these resolver-wide reinstall/ignore flags in shim mode or force dependency installation off for the kept command.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b3649d4. Resolver-wide reinstall / ignore-installed switches (pip --force-reinstall, --ignore-installed, -I; uv --reinstall) are stripped in shim mode so they cannot rebuild already-satisfied baked deps pulled in by a kept target. The kept target still installs. Added test_force_reinstall_flag_stripped / _ignore_installed_short_flag_stripped / _uv_reinstall_flag_stripped.
| # package (e.g. `uv pip install -P torch peft`) lets an ordinary install target | ||
| # refresh that package and clobber the pinned stack. Filter its value through | ||
| # _KEEP too. Unlike -e it is not itself an install target (no has_target). | ||
| _UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package"} |
There was a problem hiding this comment.
Filter uv reinstall-package selectors too
Fresh selector beyond the earlier -P fix: local uv pip install --help documents --reinstall-package <name> as reinstalling a specific package, but this protected-selector set only includes -P/--upgrade-package. As a result, uv pip install --reinstall-package=torch peft is passed through and can reinstall torch, while the separated form drops torch but leaves a dangling --reinstall-package that consumes the next kept target; route --reinstall-package values through the same _KEEP handling as upgrade-package.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in b3649d4. --reinstall-package is now in both _VALUE_FLAGS and _UPGRADE_PKG_FLAGS, so its value is filtered through _KEEP exactly like -P/--upgrade-package: a protected name drops the flag+value pair together (no dangling selector left to swallow the next target) in both --reinstall-package torch peft and --reinstall-package=torch peft. Added test_reinstall_package_protected_separated_dropped / _inline_dropped, _unprotected_kept and _transformers_pin_recorded.
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Notebook pip/uv shim (docker/unsloth_pip_shim.py), all active only under
UNSLOTH_NB_SHIM=1:
- Parse a bare wheel filename (torch-*.whl in the CWD, no ./ or / prefix) so
it is matched against _KEEP instead of passing through as an opaque
positional and reinstalling the baked torch.
- Infer the distribution from an egg-less VCS URL by repo basename
(git+https://github.com/huggingface/transformers.git -> transformers) so
the egg-less form the repo itself recommends cannot clobber the baked stack.
- Refuse remote (URL) -r/-c requirement/constraint files -- top-level and
nested includes -- since their pins cannot be inspected before the real
tool would fetch and install them.
- Strip resolver-wide reinstall/ignore-installed switches
(--force-reinstall, --ignore-installed, -I, uv --reinstall) so they cannot
rebuild already-satisfied baked deps pulled in by a kept target.
- Route uv --reinstall-package through the same _KEEP handling as
-P/--upgrade-package (both attached and separated forms; no dangling flag).
Entrypoint (docker/entrypoint.sh): select_cuda_jit_tools() now scans every
visible GPU's compute_cap instead of only the first, so a datacenter Blackwell
(sm_103/sm_121) behind an H100/B200 still enables the cu13 JIT tools it needs.
Adds regression tests for each case (tests/python/test_unsloth_pip_shim.py,
tests/sh/test_select_cuda_jit_tools.sh).
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 167fdf26b9
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| _seg = _seg.strip().lower().replace("_", "-") | ||
| if _seg: | ||
| return _seg | ||
| return None # plain url / local path -> let it pass through |
There was a problem hiding this comment.
Parse protected source archives before URL passthrough
When a notebook or requirements file names a protected package as a source archive instead of a wheel, for example pip install https://files.pythonhosted.org/.../unsloth-2026.7.1.tar.gz or ./torch-2.11.0.tar.gz, pip accepts <archive url/path> targets (checked local python -m pip install --help), but _canon() only extracts wheel basenames before returning None here. That token is kept and the real pip can reinstall unsloth/torch/transformers into the baked venv, bypassing _KEEP; parse sdist/zip basenames or reject protected archive URLs/paths before this passthrough.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6a078b1. Added a _sdist_name() helper and wired it into both the URL/path branch and the bare-filename path of _canon, so a protected source archive (https://.../unsloth-2026.7.1.tar.gz, ./torch-2.11.0.tar.gz, .zip, .tar.bz2, etc.) is matched against _KEEP instead of passing through. It splits at the first hyphen before a digit so hyphenated names like flashinfer-python survive. Added test_sdist_url/bare/zip/hyphenated_protected_dropped and _unprotected_kept.
| _REQ_FILE_FLAGS = {"-r", "--requirement"} | ||
| # Constraint files are not install targets, but pip applies their pins during | ||
| # resolution, so a `-c constraints.txt` that pins torch/transformers/etc. can | ||
| # still downgrade or reinstall a baked package when another target pulls it in. | ||
| # Filter protected packages out of them the same way as requirement files. | ||
| _CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"} |
There was a problem hiding this comment.
Route uv plural requirements through the filter
When a notebook uses uv's documented plural flags, e.g. uv pip install --requirements reqs.txt or uv pip install --constraints constraints.txt peft (checked local uv pip install --help: -r, --requirements <REQUIREMENTS> and -c, --constraints <CONSTRAINTS>), these sets only recognize the singular pip-style names. The scanner therefore forwards those files to uv unfiltered, so protected pins inside can still reinstall or constrain torch/transformers despite the requirement/constraint filtering added for the other spellings.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6a078b1. Added uv's plural long forms --requirements and --constraints to _REQ_FILE_FLAGS / _CONSTRAINT_FILE_FLAGS (and _VALUE_FLAGS), so uv pip install --requirements reqs.txt / --constraints constraints.txt go through the same protected-package filter as the singular pip-style names. Added test_uv_plural_requirements_filtered and test_uv_plural_constraints_filtered.
| # guise of installing an unprotected package. The kept target still installs; its | ||
| # already-satisfied protected deps are left untouched. Per-package selectors | ||
| # (--reinstall-package / -P) are handled through _UPGRADE_PKG_FLAGS instead. | ||
| _REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall"} |
There was a problem hiding this comment.
Strip eager upgrade flags before kept installs
For notebook cells like pip install -U --upgrade-strategy eager peft, these resolver-wide upgrade flags are still forwarded even though the local pip help says eager mode upgrades dependencies regardless of whether installed versions satisfy requirements. A kept target such as peft can then make pip upgrade protected dependencies like torch or transformers underneath the baked CUDA stack; strip -U/--upgrade with eager upgrade strategy (or otherwise force protected deps to stay installed) in the shim path too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6a078b1. --upgrade-strategy is now a value flag whose flag+value pair is dropped in shim mode (_DROP_VALUE_FLAGS), so pip install -U --upgrade-strategy eager peft falls back to pip's default only-if-needed and cannot eagerly rebuild already-satisfied baked deps like torch/transformers. The kept target (peft) still installs. Added test_upgrade_strategy_eager_dropped / _eager_inline_dropped / _only_if_needed_also_dropped.
| /opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \ | ||
| "${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do | ||
| orig="${nvrtc_dir}/libnvrtc.so.12.cu128.orig" | ||
| [[ -e "${orig}" ]] || continue |
There was a problem hiding this comment.
Do not leave CUDA 13 NVRTC active for non-root users
When the container is started with --user on a non-sm_103/sm_121 host, this symlink restore cannot write into the root-owned venv and is silently ignored. Because the Dockerfile bakes libnvrtc.so.12 pointing at CUDA 13 first, supported Ampere/Ada/Hopper/Turing hosts on 570-579 drivers keep CUDA 13 NVRTC active even though the gate decided they need cu12.8, so later NVRTC/Jiterator paths can emit cubins the accepted driver cannot load; make cu12.8 the immutable default or run the switch before dropping privileges.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6a078b1. Inverted the NVRTC default: the build now keeps cu12.8 as the immutable default (libnvrtc.so.12 -> .cu128.orig) and stages a .cu13 alias beside it; select_cuda_jit_tools retargets libnvrtc.so.12 -> .cu13 ONLY for sm_103/sm_121. Because the safe default needs no runtime write, a non-root docker run --user container -- which cannot rewrite the symlink -- now keeps cu12.8 (loadable on every 570+ driver) instead of being stuck on cu13. Applied to both the base and Studio Dockerfiles; entrypoint gating test updated to assert the cu12.8 default is left untouched for non-DC arches and switched to cu13 for sm_103/sm_121.
Notebook pip/uv shim (docker/unsloth_pip_shim.py, active only under
UNSLOTH_NB_SHIM=1):
- Parse protected source archives (sdist/zip) by basename too, e.g.
`pip install https://.../unsloth-2026.7.1.tar.gz` or `./torch-2.11.0.tar.gz`,
mirroring the wheel-basename handling. A first-hyphen-before-digit split
keeps hyphenated names like flashinfer-python intact.
- Recognise uv's PLURAL long flags --requirements / --constraints, so those
files go through the same protected-package filter as the singular names.
- Drop --upgrade-strategy eager in shim mode so a kept target cannot eagerly
rebuild already-satisfied baked deps (falls back to pip's only-if-needed).
NVRTC default (docker/Dockerfile, docker/Dockerfile.studio, docker/entrypoint.sh):
- Make cu12.8 the immutable baked default (libnvrtc.so.12 -> .cu128.orig) with
a staged .cu13 alias, and have select_cuda_jit_tools retarget to cu13 ONLY
for sm_103/sm_121. Previously cu13 was baked as the default and restored to
cu12.8 at runtime, so a non-root `docker run --user` container that cannot
rewrite the symlink stayed on cu13 NVRTC and emitted cubins a 570-579 driver
cannot load. The safe default now needs no runtime write.
Adds regression tests for each case (tests/python/test_unsloth_pip_shim.py,
tests/sh/test_select_cuda_jit_tools.sh).
|
@codex review |
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Summary
Adds a Docker setup for Unsloth that runs on any NVIDIA GPU host from Ampere through Blackwell (sm_80 through sm_120: A100, RTX 30/40, H100, B100/B200, RTX 50-series, RTX 6000 Pro Blackwell) and natively on aarch64 (GB10 / Grace, DGX Spark). Two images are published to
docker.io/unsloth/unsloth:docker/Dockerfile, tag:base): the full training stack --torch 2.10.0+cu128, Unsloth + unsloth_zoo, TRL / PEFT / accelerate, bitsandbytes, triton, xformers (amd64), vLLM -- plus llama.cpp for GGUF tooling, JupyterLab, and the baked Unsloth notebooks. Run it headless for training,unsloth-run <notebook|url>, orjupyter lab.docker/Dockerfile.studio, tag:latest): layers Unsloth Studio on top of the base image and runs the production service trio under supervisord -- Studio on 8000, JupyterLab on 8888, key-only sshd on 22 -- plus an optional Cloudflare tunnel for JupyterLab.The build itself requires no GPU at all: it runs on free GitHub-hosted runners, on a developer laptop without an NVIDIA card, or on any datacenter GPU. All produce byte-identical images.
Multi-arch. amd64 and arm64 are built in parallel on native GitHub runners (
ubuntu-latest+ubuntu-24.04-arm, both free on public repos since Aug 2025), pushed by digest, then merged into a single multi-platform manifest, sodocker pullselects the right child automatically. Native arm64 is ~3x faster than QEMU and runs on DGX Spark / Grace with CUDA working as normal (no runtime emulation).Why the build does not need a GPU
There are four places where a naive Docker build silently couples to the build-host GPU. The Dockerfile breaks each one:
1. Wheel selection. The README install line
uv pip install unsloth --torch-backend=autointrospects the build host's driver. This Dockerfile pinstorch==2.10.0against--extra-index-url https://download.pytorch.org/whl/cu128explicitly. No--torch-backend=auto, noinstall.sh.2. Dep resolution order. Splitting installs into multiple
pip installcalls letsbitsandbytes0.49.x's transitivecuda-toolkit==13.0.2dep silently upgradetorch 2.10.0+cu128 -> 2.12.0+cu130in a later pass, leaving cu128 xformers stranded. This Dockerfile collapses everything into a singleuv pip installwith--index-strategy unsafe-best-matchso the resolver sees all constraints at once.3. Build-time verification.
torch.cuda.get_arch_list()returns[]when no GPU is visible. The Dockerfile uses the raw C++ accessortorch._C._cuda_getArchFlags(), which reads compiled wheel metadata directly ('sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120'on amd64;'sm_80 sm_90 sm_100 sm_120'on aarch64). Required packages are checked viaimportlib.metadata.version()instead of importing them, becauseimport unslothtriggerstorch.cuda.get_device_properties(0), which can't be satisfied on a GPU-less host. Import-time correctness is exercised at deploy time bysmoke_test.py.4. Compiled-kernel cache. If anything imports unsloth during the build, Triton JITs kernels keyed to the build host's compute capability and bakes them into
unsloth_compiled_cache/.UNSLOTH_COMPILE_DISABLE=1andUNSLOTH_COMPILE_OVERWRITE=0prevent this. The deploy GPU produces its own cache on first use.The underlying reason all of this works is that cu128 PyTorch wheels are already fat binaries, cross-compiled upstream to include SASS for every architecture from sm_70 through sm_120 (and sm_80/90/100/120 on aarch64). The build host's GPU was never needed for the binary content, only by code that pretended to need it.
What is in the images
Base image (
:base)torch 2.10.0+cu128(held against the cu cascade), withTORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0+PTX"so any source build inside the container honours the full Ampere-through-Blackwell range. The+PTXsuffix on the top arch gives forward-compat JIT-PTX for future consumer Blackwell SKUs.unslothai/notebooksbaked in as a read-only template and synced to/workspace/unsloth-notebookson boot (edit-preserving refresh from GitHub when reachable).!pip/!uvshim that keeps the core GPU stack pinned while letting a notebook install its own extras, per-notebook transformers-version sidecars, andunsloth-runfor headless execution of any notebook or URL.[huggingface]extra (no cu128 aarch64 xformers/vLLM wheel) and installs cuda-nvrtc/nvcc from NVIDIA's sbsa repo.Full image (
:latest)arm64 / aarch64
A native arm64 child is published in the same manifest. CI covers it two ways:
docker-build-arm64-native.ymlbuilds + smoke-checks on the freeubuntu-24.04-armrunner (the path DGX Spark / Grace users take), anddocker-build-arm64-qemu.ymlcross-builds under QEMU as a fallback signal and to validate the documentedsetup_qemu.shrecipe.Validation
Base image validated end-to-end on a B200 (sm_100) host. With
CUDA_VISIBLE_DEVICES=""(simulating the GPU-less CI runner), the resolved stack istorch 2.10.0+cu128, triton, xformers (cu128 wheel from the PyTorch index), bitsandbytes, unsloth + unsloth_zoo, transformers, trl, peft, accelerate, with arch flags['sm_70','sm_75','sm_80','sm_86','sm_90','sm_100','sm_120']. Runtime path on B200 GPU 0:smoke_test.pyimports unsloth, loadsLlama-3.2-1B-Instruct-bnb-4bitin 4-bit, and completes LoRA steps with loss decreasing. The full image was launched on the same host and Studio, JupyterLab, and sshd all came up under supervisord.Files
docker/Dockerfiledocker/Dockerfile.studiodocker/entrypoint.sh,docker/supervisord.conf,docker/studio_launch.shdocker/unsloth_sync_notebooks.sh,docker/unsloth_nb_compat.py,docker/unsloth_nb_content_sig.pydocker/unsloth_pip_shim.py,docker/unsloth_run.py,docker/unsloth_ipython_startup.py!pip/!uvshim, headless runner, IPython startup hookdocker/unsloth_studio_update.sh,docker/unsloth_llama_update.sh,docker/unsloth_jupyter_tunnel.sh,docker/fetch_llama_prebuilt.pydocker/run.sh,docker/build.sh,docker/test_locally.sh,docker/docker_confirm.sh,docker/docker_confirm.ps1,docker/setup_qemu.shdocker/freeze.sh,docker/hf_pull.sh,docker/hf_push.sh,docker/smoke_test.py,docker/.dockerignore.github/workflows/docker-publish.ymlDesign notes
nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04for the build stage,-cudnn-runtime-ubuntu24.04for deploy. No nvcc in the published amd64 image./opt/unsloth-venv/requirements.lock.txtinside the image and can be extracted withdocker/freeze.shfor a fully-pinned rebuild later.Test plan
DOCKERHUB_USERNAMEandDOCKERHUB_TOKENunsloth/unsloth:lateston an RTX 50-series host and confirm Studio (8000) + JupyterLab (8888) come upvars.HAS_GPU_RUNNER=trueonce a self-hosted GPU runner is registered so the post-publish smoke test exercises real sm_120 paths